#!/usr/bin/python3

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import logging
from math import pi

from torch import device as torch_device
from torch import full as torch_full
from torch import Tensor, chunk, index_select, cos, sin, stack, no_grad, argsort
from torch import zeros as torch_zeros
from torch import add as torch_add
from torch import norm as torch_norm
from torch import mul as torch_mul
from torch import log as torch_log
from torch import div as torch_div
from torch import abs as torch_abs
from torch.nn import Module, Parameter, Linear
from torch.nn.init import uniform_
from torch.nn.functional import softmax, logsigmoid


class KGEModel(Module):
    def __init__(self, model_name, nentity, nrelation, hidden_dim, gamma, use_cuda=False, max_head_len=4,
                 max_body_len=1, use_delta=False, init_delta=0.5, delta_rate=0.1, double_entity_embedding=False, double_relation_embedding=False):
        super(KGEModel, self).__init__()
        self.model_name = model_name
        self.nentity = nentity
        self.nrelation = nrelation
        self.hidden_dim = hidden_dim
        self.epsilon = 2.0
        self.use_cuda = use_cuda
        self.device = torch_device("cuda") if self.use_cuda else torch_device("cpu")
        self.max_head_len = max_head_len
        self.max_body_len = max_body_len
        self.delta_rate = delta_rate
        self.init_delta = init_delta

        self.gamma = Parameter(
            Tensor([gamma]),
            requires_grad=False
        )

        self.embedding_range = Parameter(
            Tensor([(self.gamma.item() + self.epsilon) / hidden_dim]),
            requires_grad=False
        )

        if use_delta:
            self.delta_h = Parameter(torch_full((self.nentity, self.nrelation), init_delta), requires_grad=True)
            self.delta_t = Parameter(torch_full((self.nentity, self.nrelation), init_delta), requires_grad=True)
        else:
            self.delta_h = Parameter(torch_full((self.nentity, self.nrelation), init_delta), requires_grad=False)
            self.delta_t = Parameter(torch_full((self.nentity, self.nrelation), init_delta), requires_grad=False)

        self.entity_dim = hidden_dim * 2 if double_entity_embedding else hidden_dim
        self.relation_dim = hidden_dim * 2 if double_relation_embedding else hidden_dim

        self.entity_embedding = Parameter(torch_zeros(nentity, self.entity_dim))
        uniform_(
            tensor=self.entity_embedding,
            a=-self.embedding_range.item(),
            b=self.embedding_range.item()
        )

        self.relation_embedding = Parameter(torch_zeros(nrelation, self.relation_dim))
        uniform_(
            tensor=self.relation_embedding,
            a=-self.embedding_range.item(),
            b=self.embedding_range.item()
        )
        if model_name == "TransH":
            # self.relation_project_matrix = nn.Parameter(torch.ones())
            pass

        if model_name == 'pRotatE':
            self.modulus = Parameter(Tensor([[0.5 * self.embedding_range.item()]]))

        # Do not forget to modify this line when you add a new model in the "forward" function
        if model_name not in ['TransE', 'DistMult', 'ComplEx', 'RotatE', 'pRotatE']:
            raise ValueError('model %s not supported' % model_name)

        if model_name == 'RotatE' and (not double_entity_embedding or double_relation_embedding):
            raise ValueError('RotatE should use --double_entity_embedding')

        if model_name == 'ComplEx' and (not double_entity_embedding or not double_relation_embedding):
            raise ValueError('ComplEx should use --double_entity_embedding and --double_relation_embedding')

        self.generate_delta = Linear((2*self.entity_dim + self.relation_dim), 1)


    def forward(self, sample, mode='single', use_delta=False, use_pattern=False):
        '''
        Forward function that calculate the score of a batch of triples.
        In the 'single' mode, sample is a batch of triple.
        In the 'head-batch' or 'tail-batch' mode, sample consists two part.
        The first part is usually the positive sample.
        And the second part is the entities in the negative samples.
        Because negative samples and positive samples usually share two elements
        in their triple ((head, relation) or (relation, tail)).
        '''

        if mode == 'single':
            batch_size, negative_sample_size = sample.size(0), 1

            head = index_select(
                self.entity_embedding,
                dim=0,
                index=sample[:, 0]
            ).unsqueeze(1)

            relation = index_select(
                self.relation_embedding,
                dim=0,
                index=sample[:, 1]
            ).unsqueeze(1)

            tail = index_select(
                self.entity_embedding,
                dim=0,
                index=sample[:, 2]
            ).unsqueeze(1)

            if use_delta:
                delta_h_index, delta_r_index, delta_t_index = sample[:, 0], sample[:, 1], sample[:, 2]
                delta_h = self.delta_h[delta_h_index, delta_r_index]
                delta_t = self.delta_t[delta_t_index, delta_r_index]
                deltas = torch_add(delta_h, delta_t)
                # deltas = self.generate_delta(torch.cat((head, relation, tail), -1))
                if self.use_cuda:
                    deltas = deltas.cuda()
                data = (head, relation, tail, deltas, mode)
            else:
                data = (head, relation, tail, mode)

        elif mode == 'head-batch':
            tail_part, head_part = sample
            batch_size, negative_sample_size = head_part.size(0), head_part.size(1)

            head = index_select(
                self.entity_embedding,
                dim=0,
                index=head_part.view(-1)
            ).view(batch_size, negative_sample_size, -1)

            relation = index_select(
                self.relation_embedding,
                dim=0,
                index=tail_part[:, 1]
            ).unsqueeze(1)

            tail = index_select(
                self.entity_embedding,
                dim=0,
                index=tail_part[:, 2]
            ).unsqueeze(1)

            if use_delta:
                delta_h_index = head_part
                delta_r_index = tail_part[:, 1].unsqueeze(1).repeat(1, negative_sample_size).view(batch_size, negative_sample_size)
                delta_t_index = tail_part[:, 2].unsqueeze(1).repeat(1, negative_sample_size).view(batch_size, negative_sample_size)
                delta_h = self.delta_h[delta_h_index, delta_r_index]
                delta_t = self.delta_t[delta_t_index, delta_r_index]
                deltas = torch_add(delta_h, delta_t)
                # deltas = self.generate_delta(torch.cat((head, relation, tail), -1))
                if self.use_cuda:
                    deltas = deltas.cuda()
                data = (head, relation, tail, deltas, mode)
            else:
                data = (head, relation, tail, mode)

        elif mode == 'tail-batch':
            head_part, tail_part = sample
            batch_size, negative_sample_size = tail_part.size(0), tail_part.size(1)

            head = index_select(
                self.entity_embedding,
                dim=0,
                index=head_part[:, 0]
            ).unsqueeze(1)

            relation = index_select(
                self.relation_embedding,
                dim=0,
                index=head_part[:, 1]
            ).unsqueeze(1)

            tail = index_select(
                self.entity_embedding,
                dim=0,
                index=tail_part.view(-1)
            ).view(batch_size, negative_sample_size, -1)

            if use_delta:
                delta_h_index = head_part[:, 0].unsqueeze(1).repeat(1, negative_sample_size).view(batch_size, negative_sample_size)
                delta_r_index = head_part[:, 1].unsqueeze(1).repeat(1, negative_sample_size).view(batch_size, negative_sample_size)
                delta_t_index = tail_part
                delta_h = self.delta_h[delta_h_index, delta_r_index]
                delta_t = self.delta_t[delta_t_index, delta_r_index]
                deltas = torch_add(delta_h, delta_t)
                # deltas = self.generate_delta(torch.cat((head, relation, tail), -1))
                if self.use_cuda:
                    deltas = deltas.cuda()
                data = (head, relation, tail, deltas, mode)
            else:
                data = (head, relation, tail, mode)

        else:
            raise ValueError('mode %s not supported' % mode)

        model_func = {
            'TransE': self.TransE,
            'DistMult': self.DistMult,
            'ComplEx': self.ComplEx,
            'RotatE': self.RotatE,
            'pRotatE': self.pRotatE
        }

        if self.model_name in model_func:
            score = model_func[self.model_name](data, mode, use_delta)
            # if mode == "single" and use_delta:
            #     score -= torch.pow(delta_h, 2) * self.delta_rate + torch.pow(delta_t, 2) * self.delta_rate
        else:
            raise ValueError('model %s not supported' % self.model_name)

        return score

    def TransE(self, samples, mode, use_delta=False):
        if use_delta:
            head, relation, tail, deltas, mode = samples
        else:
            head, relation, tail, mode = samples

        if mode == 'head-batch':
            score = head + (relation - tail)
        else:
            score = (head + relation) - tail

        score = torch_norm(score, p=1, dim=2)

        if use_delta:
            square_delta = torch_mul(deltas, deltas)
            score = torch_div(score, square_delta * 2) + torch_log(deltas)

            if mode == "single":
                score += 0.1 * square_delta
            else:
                score += 1. * (1 / square_delta)

            if self.use_cuda:
                score += torch_log(Tensor([2 * pi])).cuda()
            else:
                score += torch_log(Tensor([2 * pi]))

        score = self.gamma.item() - score
        return score

    def TransH(self, head, relation, tail, mode):
        pass

    def DistMult(self, samples, mode, use_delta=False):
        if use_delta:
            head, relation, tail, deltas, mode = samples
        else:
            head, relation, tail, mode = samples

        if mode == 'head-batch':
            score = head * (relation * tail)
        else:
            score = (head * relation) * tail

        score = score.sum(dim=2)
        return score

    def ComplEx(self, samples, mode, use_delta=False):
        if use_delta:
            head, relation, tail, deltas, mode = samples
        else:
            head, relation, tail, mode = samples

        re_head, im_head = chunk(head, 2, dim=2)
        re_relation, im_relation = chunk(relation, 2, dim=2)
        re_tail, im_tail = chunk(tail, 2, dim=2)

        if mode == 'head-batch':
            re_score = re_relation * re_tail + im_relation * im_tail
            im_score = re_relation * im_tail - im_relation * re_tail
            score = re_head * re_score + im_head * im_score
        else:
            re_score = re_head * re_relation - im_head * im_relation
            im_score = re_head * im_relation + im_head * re_relation
            score = re_score * re_tail + im_score * im_tail

        score = score.sum(dim=2)
        return score

    def RotatE(self, samples, mode, use_delta=False):
        pi = 3.14159265358979323846
        if use_delta:
            head, relation, tail, deltas, mode = samples
        else:
            head, relation, tail, mode = samples

        re_head, im_head = chunk(head, 2, dim=2)
        re_tail, im_tail = chunk(tail, 2, dim=2)

        phase_relation = relation / (self.embedding_range.item() / pi)

        re_relation = cos(phase_relation)
        im_relation = sin(phase_relation)

        if mode == 'head-batch':
            re_score = re_relation * re_tail + im_relation * im_tail
            im_score = re_relation * im_tail - im_relation * re_tail
            re_score = re_score - re_head
            im_score = im_score - im_head
        else:
            re_score = re_head * re_relation - im_head * im_relation
            im_score = re_head * im_relation + im_head * re_relation
            re_score = re_score - re_tail
            im_score = im_score - im_tail

        score = stack([re_score, im_score], dim=0)
        score = score.norm(dim=0).sum(dim=2)

        if use_delta:
            square_delta = torch_mul(deltas, deltas)
            score = torch_div(score, square_delta) + torch_log(deltas)

        score = self.gamma.item() - score
        return score

    def pRotatE(self, samples, mode, use_delta=False):
        if use_delta:
            head, relation, tail, deltas, mode = samples
        else:
            head, relation, tail, mode = samples

        pi = 3.14159262358979323846

        # Make phases of entities and relations uniformly distributed in [-pi, pi]

        phase_head = head / (self.embedding_range.item() / pi)
        phase_relation = relation / (self.embedding_range.item() / pi)
        phase_tail = tail / (self.embedding_range.item() / pi)

        if mode == 'head-batch':
            score = phase_head + (phase_relation - phase_tail)
        else:
            score = (phase_head + phase_relation) - phase_tail

        score = sin(score)
        score = torch_abs(score)

        score = self.gamma.item() - score.sum(dim=2) * self.modulus
        return score

    @staticmethod
    def train_step(model, optimizer, train_iterator, args):
        '''
        A single train step. Apply back-propation and return the loss
        '''

        model.train()

        optimizer.zero_grad()

        positive_sample, negative_sample, subsampling_weight, mode = next(train_iterator)

        if args.cuda:
            positive_sample = positive_sample.cuda()
            negative_sample = negative_sample.cuda()
            subsampling_weight = subsampling_weight.cuda()

        negative_score = model((positive_sample, negative_sample), mode=mode, use_delta=args.use_delta)
        # negative_sample_loss = negative_score.mean(dim=1)

        if args.negative_adversarial_sampling:
            # In self-adversarial sampling, we do not apply back-propagation on the sampling weight
            negative_score = (softmax(negative_score * args.adversarial_temperature, dim=1).detach()
                              * logsigmoid(-negative_score)).sum(dim=1)
        else:
            negative_score = logsigmoid(-negative_score).mean(dim=1)

        positive_score = model(positive_sample, use_delta=args.use_delta)
        # positive_sample_loss = positive_score.mean(dim=1)

        positive_score = logsigmoid(positive_score)

        if args.uni_weight:
            positive_sample_loss = - positive_score.mean()
            negative_sample_loss = - negative_score.mean()

        else:
            positive_sample_loss = - (subsampling_weight * positive_score).sum() / subsampling_weight.sum()
            negative_sample_loss = - (subsampling_weight * negative_score).sum() / subsampling_weight.sum()

        loss = (positive_sample_loss + negative_sample_loss) / 2

        if args.regularization != 0.0:
            # Use L3 regularization for ComplEx and DistMult
            regularization = args.regularization * (
                    model.entity_embedding.norm(p=3) ** 3 +
                    model.relation_embedding.norm(p=3).norm(p=3) ** 3
            )
            loss = loss + regularization
            regularization_log = {'regularization': regularization.item()}
        else:
            regularization_log = {}

        loss.backward()

        optimizer.step()

        log = {
            **regularization_log,
            'positive_sample_loss': positive_sample_loss.item(),
            'negative_sample_loss': negative_sample_loss.item(),
            'loss': loss.item()
        }

        return log

    @staticmethod
    def test_step(model, test_dataset_list, args):
        '''
        Evaluate the model on test or valid datasets
        '''

        model.eval()
        # Otherwise use standard (filtered) MRR, MR, HITS@1, HITS@3, and HITS@10 metrics
        # Prepare dataloader for evaluation


        logs = []

        step = 0
        total_steps = sum([len(dataset) for dataset in test_dataset_list])

        with no_grad():
            for test_dataset in test_dataset_list:
                for positive_sample, negative_sample, filter_bias, mode in test_dataset:
                    if args.cuda:
                        positive_sample = positive_sample.cuda()
                        negative_sample = negative_sample.cuda()
                        filter_bias = filter_bias.cuda()

                    batch_size = positive_sample.size(0)

                    score = model((positive_sample, negative_sample), mode)
                    score += filter_bias

                    # Explicitly sort all the entities to ensure that there is no test exposure bias
                    argsort = argsort(score, dim=1, descending=True)

                    if mode == 'head-batch':
                        positive_arg = positive_sample[:, 0]
                    elif mode == 'tail-batch':
                        positive_arg = positive_sample[:, 2]
                    else:
                        raise ValueError('mode %s not supported' % mode)

                    for i in range(batch_size):
                        # Notice that argsort is not ranking
                        ranking = (argsort[i, :] == positive_arg[i]).nonzero()
                        assert ranking.size(0) == 1

                        # ranking + 1 is the true ranking used in evaluation metrics
                        ranking = 1 + ranking.item()
                        logs.append({
                            'MRR': 1.0 / ranking,
                            'MR': float(ranking),
                            'HITS@1': 1.0 if ranking <= 1 else 0.0,
                            'HITS@3': 1.0 if ranking <= 3 else 0.0,
                            'HITS@10': 1.0 if ranking <= 10 else 0.0,
                        })

                    if step % args.test_log_steps == 0:
                        logging.info('Evaluating the model... (%d/%d)' % (step, total_steps))

                    step += 1

        metrics = {}
        for metric in logs[0].keys():
            metrics[metric] = sum([log[metric] for log in logs]) / len(logs)

        return metrics
